/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.yuanrong.jobexecutor;

import com.yuanrong.exception.YRException;
import com.yuanrong.runtime.util.Constants;
import com.yuanrong.runtime.util.Utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * The JobExecutor actor runs in remote runtime to manage driver code.
 *
 * @since 2023 /06/06
 */
public class JobExecutor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JobExecutor.class);

    /**
     * The list of programs that are allowed to run on the processBuilder.
     */
    private static final List<String> ALLOWED_PROGRAM = Collections.singletonList("python3.9");

    /**
     * The key to yuanrong functionProxy port in remote JobExecutor runtime
     * environment parameter map. The value of the port will be pass to the
     * attached-runtime process as it' environment parameter.
     */
    private static final String DRIVER_SERVER_PORT = "DRIVER_SERVER_PORT";

    /**
     * The key to downloaded driver code path in JobExcutor actor runtime envronment.
     */
    private static final String ENV_DELEGATE_DOWNLOAD = "ENV_DELEGATE_DOWNLOAD";

    /**
     * The key to JobExcutor actor instance ID, which is also called 'userJobID'
     * in the context of JobExecutor.
     */
    private static final String INSTANCE_ID = "INSTANCE_ID";
    private static final String LOG_PREFIX = "(JobExecutor) ";
    private static final String DATASYSTEM_ADDR = "DATASYSTEM_ADDR";
    private static final String HOST_IP = "HOST_IP";
    private static final String YR_SPARK_HOME = "YR_SPARK_HOME";
    private static final String YR_SPARK_FUNC_URN = "YR_SPARK_FUNC_URN";
    private static final String POSIX_LISTEN_ADDR = "POSIX_LISTEN_ADDR";
    private static final String JAVA_HOME = "JAVA_HOME";
    private static final String YR_JOB_AFFINITY = "YR_JOB_AFFINITY";
    private static final String DEFAULT_DS_ADDRESS_PORT = "31501";
    private static final String DEFAULT_DRIVER_SERVER_PORT = "22771";
    private static final String DEFAULT_JAVA_HOME = "/opt/function/runtime/java8/rtsp/jre";
    private static final String DEFAULT_SPARK_HOME = "/dcache/layer/func/bucket-jobexecutor-test/spark_test.zip";
    private static final String DEFAULT_YR_SPARK_FUNC_URN =
    "sn:cn:yrk:12345678901234561234567890123456:function:0-sparkonyr-core:$latest";
    private static final int CODE_PATH_INDEX = 1;

    /**
     * The thread pool for asynchronous log printing and waiting
     * for the end of the attached-runtime.
     */
    private ThreadPoolExecutor attachedRuntimeThreadPool = new ThreadPoolExecutor(3, 3, 60L,
            TimeUnit.SECONDS, new SynchronousQueue<>());
    private int attachedRuntimePid = -1;
    private Process attachedRuntime;
    private String jobName = "";
    private String userJobID;
    private LocalDateTime jobStartTime;
    private volatile LocalDateTime jobEndTime;
    private volatile String errorMessage;
    private RuntimeEnv runtimeEnv;
    private volatile YRJobStatus attachedRuntimeStatus = YRJobStatus.RUNNING;
    private int exitCode;
    private BufferedReader inputReader;
    private BufferedReader errReader;

    /**
     * The Constructor of JobExecutor.
     *
     * @param jobName         the user-defined job name.
     * @param runtimeEnv      the python environment to be installed in remote
     *                        runtime.
     * @param localEntryPoint the entry point to be executed in remote runtime. it
     *                        will contains the localCodePath if provided by user.
     * @param affinityMsgJsonStr the affinityMsgJsonStr
     * @throws YRException the actor task exception.
     */
    public JobExecutor(String jobName, RuntimeEnv runtimeEnv, ArrayList<String> localEntryPoint,
            String affinityMsgJsonStr) throws YRException {
        this.runtimeEnv = runtimeEnv;
        this.jobName = jobName;
        this.userJobID = Utils.getEnvWithDefualtValue(INSTANCE_ID, "", LOG_PREFIX);
        LOGGER.debug("(JobExecutor) Job submitted, jobExecutor actor instanceID: {}", userJobID);

        String downloadCodePath = System.getenv(ENV_DELEGATE_DOWNLOAD);
        if (downloadCodePath != null && !downloadCodePath.isEmpty()) {
            localEntryPoint.set(CODE_PATH_INDEX,
                    String.join(Constants.BACKSLASH, downloadCodePath, localEntryPoint.get(CODE_PATH_INDEX)));
        }

        LOGGER.info("(JobExecutor) Starts attached-runtime process, entrypoint: {}", localEntryPoint);
        startAttachedRuntime(generateAttachedRuntimeEnv(affinityMsgJsonStr), localEntryPoint);
        asyncReadProcessStream();
        asyncWaitForAttachedRuntime();
    }

    /**
     * Gets the jobInfo of the the attached-runtime.
     *
     * @param isWithStatic whether invoke static infomation from remote JobExecutor
     *                     runtime.
     * @return YRJobInfo
     */
    public YRJobInfo getJobInfo(boolean isWithStatic) {
        YRJobInfo jobInfo = new YRJobInfo();
        if (isWithStatic) {
            jobInfo.setUserJobID(this.userJobID);
            jobInfo.setJobName(this.jobName);
            jobInfo.setRuntimeEnv(this.runtimeEnv);
            if (this.jobStartTime != null) {
                jobInfo.setJobStartTime(this.jobStartTime.toString());
            }
        }
        jobInfo.setStatus(this.attachedRuntimeStatus);
        jobInfo.setErrorMessage(this.errorMessage);
        if (this.jobEndTime != null) {
            jobInfo.setJobEndTime(this.jobEndTime.toString());
        }
        return jobInfo;
    }

    /**
     * Terminates the attached-runtime process.
     * This method does NOT immediately free the resources occupied by the
     * attached-runtime process.
     */
    public void stop() {
        // If 'stop()' is invoked more than once,
        // 'jobEndTime' should not be changed.
        // Therefore, an 'if' judgment is required here.
        if (this.jobEndTime == null) {
            this.jobEndTime = LocalDateTime.now();
        }
        writeStatus(YRJobStatus.STOPPED);
        if (this.attachedRuntime != null) {
            this.attachedRuntime.destroy();
        }
        this.clearThreadPool();
        LOGGER.info("(JobExecutor) Stop the attached-runtime process (pid: {}), job name: {}."
                + " Current job status: {}, stop time: {}.", this.attachedRuntimePid, this.jobName,
                this.attachedRuntimeStatus, this.jobEndTime);
    }

    private Map<String, String> generateAttachedRuntimeEnv(String affinityMsgJsonStr) throws YRException {
        String defaultDsAddr = System.getenv(HOST_IP) + ":" + DEFAULT_DS_ADDRESS_PORT;

        Map<String, String> arEnv = new HashMap<String, String>();
        arEnv.put(YR_SPARK_HOME, Utils.getEnvWithDefualtValue(YR_SPARK_HOME, DEFAULT_SPARK_HOME, LOG_PREFIX));
        arEnv.put(YR_SPARK_FUNC_URN,
                Utils.getEnvWithDefualtValue(YR_SPARK_FUNC_URN, DEFAULT_YR_SPARK_FUNC_URN, LOG_PREFIX));
        arEnv.put(DATASYSTEM_ADDR, Utils.getEnvWithDefualtValue(DATASYSTEM_ADDR, defaultDsAddr, LOG_PREFIX));
        arEnv.put(POSIX_LISTEN_ADDR, Utils.getEnvWithDefualtValue(POSIX_LISTEN_ADDR, "", LOG_PREFIX));
        arEnv.put(DRIVER_SERVER_PORT,
                Utils.getEnvWithDefualtValue(DRIVER_SERVER_PORT, DEFAULT_DRIVER_SERVER_PORT, LOG_PREFIX));
        arEnv.put(JAVA_HOME, Utils.getEnvWithDefualtValue(JAVA_HOME, DEFAULT_JAVA_HOME, LOG_PREFIX));
        arEnv.put(Constants.AUTHORIZATION, Utils.getEnvWithDefualtValue(Constants.AUTHORIZATION, "", LOG_PREFIX));

        String runtimeEnvCommand = this.runtimeEnv.toCommand();
        arEnv.put(Constants.POST_START_EXEC, runtimeEnvCommand);
        LOGGER.debug("(JobExecutor) Sets environment key-value pair({}: {}) for job (jobName: {}).",
        Constants.POST_START_EXEC, runtimeEnvCommand, this.jobName);

        arEnv.put(YR_JOB_AFFINITY, affinityMsgJsonStr);
        LOGGER.debug("(JobExecutor) Sets environment key-value pair({}: {}) for job (jobName: {}).",
                YR_JOB_AFFINITY, affinityMsgJsonStr, this.jobName);

        return arEnv;
    }

    private void startAttachedRuntime(Map<String, String> attachedRuntimeEnv, List<String> entryPoint) {
        String program = entryPoint.get(0);
        if (!ALLOWED_PROGRAM.contains(program)) {
            LOGGER.error("(JobExecutor) Program {} is not allowed to be executed for the security reason. "
                    + "Allowed programs are: {}", program, ALLOWED_PROGRAM);
            writeStatus(YRJobStatus.FAILED);
            return;
        }
        ProcessBuilder processBuilder = new ProcessBuilder(entryPoint);
        Map<String, String> env = processBuilder.environment();
        env.putAll(attachedRuntimeEnv);

        try {
            this.attachedRuntime = processBuilder.start();
            this.jobStartTime = LocalDateTime.now();
        } catch (IOException e) {
            LOGGER.error("(JobExecutor) Failed to start the attached-runtime process for the job: {}. Exception: ",
                    this.jobName, e);
            this.errorMessage = "Failed to start the attached-runtime process, userJobID: " + this.userJobID;
            writeStatus(YRJobStatus.FAILED);
            return;
        }

        this.attachedRuntimePid = getProcessPid(this.attachedRuntime);
        LOGGER.info(
                "(JobExecutor) Succeeded to start a process as an attached-runtime(pid: {}) to run driver code: {}",
                this.attachedRuntimePid, this.jobName);
    }

    private void asyncWaitForAttachedRuntime() {
        if (this.attachedRuntime == null) {
            LOGGER.warn("(JobExecutor) Attached-runtime has not been started.");
            return;
        }
        this.attachedRuntimeThreadPool.execute(() -> {
            try {
                this.exitCode = this.attachedRuntime.waitFor();
                if (this.jobEndTime == null) {
                    this.jobEndTime = LocalDateTime.now();
                }
            } catch (InterruptedException e) {
                if (this.jobEndTime != null) {
                    // If 'this.jobEndTime' has not null value,
                    // attached-runtime is stopped by calling 'stop()',
                    // which is not regarded as an unexpected error.
                    return;
                }
                LOGGER.warn("(JobExecutor) Attached-runtime process (pid: {}) exits abnormally, Exception: ",
                        this.attachedRuntimePid, e);
                writeStatus(YRJobStatus.FAILED);
                this.clearThreadPool();
                return;
            }
            LOGGER.info("(JobExecutor) Attached-runtime process(pid: {}) exit. Exit code: {}",
                    this.attachedRuntimePid, this.exitCode);
            if (this.exitCode != 0) {
                writeStatus(YRJobStatus.FAILED);
                this.clearThreadPool();
                return;
            }
            writeStatus(YRJobStatus.SUCCEEDED);
            /*
             * There is a one-to-one correspondence between JobExecutor runtime and
             * attached-runtime. Therefore, after attached-runtime is finished, the thread pool belonging
             * to JobExecutor runtime can be closed.
             */
            this.clearThreadPool();
        });
    }

    private int getProcessPid(Process process) {
        int pid = -1;
        Field field;
        try {
            field = process.getClass().getDeclaredField("pid");
            field.setAccessible(true);
            pid = (int) field.get(process);
        } catch (NoSuchFieldException | SecurityException | IllegalArgumentException | IllegalAccessException e) {
            LOGGER.warn("(JobExecutor) Failed to get pid of driver code (jobName: {}) process. Exception: ",
                    this.jobName, e);
        }
        return pid;
    }

    private void writeStatus(YRJobStatus status) {
        switch (status) {
            case RUNNING:
                this.attachedRuntimeStatus = status;
                break;
            default:
                if (this.attachedRuntimeStatus == null || this.attachedRuntimeStatus == YRJobStatus.RUNNING) {
                    this.attachedRuntimeStatus = status;
                }
        }
    }

    private void asyncReadProcessStream() {
        if (this.attachedRuntime == null) {
            LOGGER.warn("(JobExecutor) Attached-runtime has not been started.");
            return;
        }

        this.errReader = new BufferedReader(
            new InputStreamReader(this.attachedRuntime.getErrorStream(), StandardCharsets.UTF_8));
        this.attachedRuntimeThreadPool.execute(() -> {
            String line;
            this.errorMessage = "";
            try {
                while ((line = errReader.readLine()) != null) {
                    this.errorMessage = this.errorMessage + System.lineSeparator() + line;
                }
            } catch (IOException e) {
                LOGGER.error("(JobExecutor) Failed to read attached-runtime error message, Exception: ", e);
                this.errorMessage = "Failed to read attached-runtime error message. Exception: " + e.getMessage();
            }
            if (!this.errorMessage.isEmpty()) {
                LOGGER.error("(Attached-runtime) {}", this.errorMessage);
            }
        });

        this.inputReader = new BufferedReader(
            new InputStreamReader(this.attachedRuntime.getInputStream(), StandardCharsets.UTF_8));
        this.attachedRuntimeThreadPool.execute(() -> {
            String line;
            try {
                while ((line = inputReader.readLine()) != null) {
                    LOGGER.info("(Attached-runtime) {}", line);
                }
            } catch (IOException e) {
                LOGGER.error("(JobExecutor) Failed to read attached-runtime inputStream message. Exception: ", e);
            }
        });
    }

    private void clearThreadPool() {
        try {
            if (this.inputReader != null) {
                this.inputReader.close();
            }
        } catch (IOException e) {
            LOGGER.error("(JobExecutor) Failed to close attached-runtime input stream, Exception: ", e);
        }

        try {
            if (errReader != null) {
                errReader.close();
            }
        } catch (IOException e) {
            LOGGER.error("(JobExecutor) Failed to close attached-runtime error stream, Exception: ", e);
        }

        this.attachedRuntimeThreadPool.shutdownNow();
    }
}
